﻿#nullable enable

namespace NServiceBus;

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using MessageMutator;
using Microsoft.Extensions.DependencyInjection;
using Pipeline;

[System.Diagnostics.CodeAnalysis.SuppressMessage("Code", "PS0025:Dictionary keys should implement GetHashCode", Justification = "Mutators are registered based on reference equality")]
class MutateIncomingMessageBehavior(HashSet<IMutateIncomingMessages> mutators)
    : IBehavior<IIncomingLogicalMessageContext, IIncomingLogicalMessageContext>
{
    public Task Invoke(IIncomingLogicalMessageContext context, Func<IIncomingLogicalMessageContext, Task> next) => hasIncomingMessageMutators ? InvokeIncomingMessageMutators(context, next) : next(context);

    async Task InvokeIncomingMessageMutators(IIncomingLogicalMessageContext context, Func<IIncomingLogicalMessageContext, Task> next)
    {
        var logicalMessage = context.Message;
        var current = logicalMessage.Instance;

        var mutatorContext = new MutateIncomingMessageContext(current, context.Headers, context.CancellationToken);

        var hasMutators = false;

        foreach (var mutator in context.Builder.GetServices<IMutateIncomingMessages>())
        {
            hasMutators = true;

            await mutator.MutateIncoming(mutatorContext)
                .ThrowIfNull()
                .ConfigureAwait(false);
        }

        foreach (var mutator in mutators)
        {
            hasMutators = true;

            await mutator.MutateIncoming(mutatorContext)
                .ThrowIfNull()
                .ConfigureAwait(false);
        }

        hasIncomingMessageMutators = hasMutators;

        if (mutatorContext.MessageInstanceChanged)
        {
            context.UpdateMessageInstance(mutatorContext.Message);
        }

        await next(context).ConfigureAwait(false);
    }

    volatile bool hasIncomingMessageMutators = true;
}